"""
@Time: 2021/1/24 下午 7:12
@Author: jinzhuan
@File: et_toolkit.py
@Desc: 
"""
import torch
from cognlp import *

torch.cuda.set_device(5)

tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run(
    "She would be the first foreign woman to die in the wave of kidnappings in Iraq.")

# ner_toolkit = NerToolkit(
#     bert_model='bert-base-cased',
#     model_path='../../../cognlp/data/ner/trex/toolkit/model.pkl',
#     vocabulary_path='../../../cognlp/data/ner/trex/toolkit/vocabulary.txt',
#     device=torch.device("cuda"),
#     device_ids=[5],
#     max_seq_length=256)
ner_toolkit = NerAce2005Toolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/ner/ace2005/toolkit/model.pt',
    vocabulary_path='../../../cognlp/data/ner/ace2005/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)
ner_result = ner_toolkit.run(words)
print(ner_result)

et_toolkit = EtToolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/et/OntoNotes/toolkit/model.pkl',
    vocabulary_path='../../../cognlp/data/et/OntoNotes/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)
et_result = et_toolkit.run(words, ner_result)

print(et_result)
